#!/usr/bin env python3
# -*- coding:UTF8 -*-

"""
"""
import cv2 as cv
import numpy as np
import os

from ai_library.components.splitPlate import splitContour

bin_n = 16

# 来自opencv的sample，用于svm训练
def preprocess_hog(digits):
    samples = []
    for img in digits:
        samples.append(HOG(img))
    return np.float32(samples)


def HOG(img):
    winSize = (20, 20)
    blockSize = (10, 10)
    blockStride = (5, 5)
    cellSize = (10, 10)
    nbins = 9
    derivAperture = 1
    winSigma = -1.
    histogramNormType = 0
    L2HysThreshold = 0.2
    gammaCorrection = 1
    nlevels = 64
    signedGradients = True

    hog = cv.HOGDescriptor(winSize,
                           blockSize,
                           blockStride,
                           cellSize,
                           nbins,
                           derivAperture,
                           winSigma,
                           histogramNormType,
                           L2HysThreshold,
                           gammaCorrection,
                           nlevels,
                           signedGradients)
    return hog.compute(img)


class StatModel(object):
    def load(self, fn):
        self.model = self.model.load(fn)

    def save(self, fn):
        self.model.save(fn)


class SVM(StatModel):
    def __init__(self, C = .5, gamma = .5):
        self.model = cv.ml.SVM_create()
        self.model.setGamma(gamma)
        self.model.setC(C)
        self.model.setKernel(cv.ml.SVM_LINEAR)
        self.model.setType(cv.ml.SVM_C_SVC)

    # 训练svm
    def train(self, samples, responses):
        self.model.train(samples, cv.ml.ROW_SAMPLE, responses)

    # 字符识别
    def predict(self, samples):
        r = self.model.predict(samples)
        return r[1].ravel()

model = SVM()

def train_svm():
    global model
    dirpath = os.path.dirname(__file__)
    svm_path = os.path.join(dirpath, "./ai_library/models/svm.xml")
    print(svm_path)
    if os.path.exists(svm_path):
        model.load('./ai_library/svm.xml')
    else:
        chars_train = []
        chars_label = []
        for root, dirs, files in os.walk(os.path.join(dirpath, "/train/chars2")):
            if len(os.path.basename(root)) > 1:
                continue
            root_int = ord(os.path.basename(root))
            for filename in files:
                filepath = os.path.join(root, filename)
                gray = cv.imread(filepath, 0)
                chars_train.append(gray)
                chars_label.append(root_int)
        train_data = preprocess_hog(chars_train)
        train_label = np.array(chars_label)
        train_label = train_label[:, np.newaxis]
        print(train_label.shape)
        model.train(train_data, train_label)
        model.save(svm_path)
        print('训练完成')
    return model


def test(frame):
    model = train_svm()
    samples = []
    samples.append(splitContour(frame))
    if not samples is None:
        for s in samples:
            res = []
            for k in sorted(s.keys(), key = int)[1:]:
                hog = HOG(s[k])
                hog = np.float32(hog)
                hog = hog.reshape(1, 81)
                res.append(chr(model.predict(hog)))
            plate = ''.join(res)
            return plate


if __name__ == '__main__':
    train_svm()
